from tables.vector import Vectors
from utils.main import get_keywords, get_embedding, get_rerank_scores


async def baseline_search_chunks(
        query: str,
        llm: str,
        embed: str,
        collection_name: str,
        retrieval_type: int,
        reranker: str,
        content_augmentation:int,
        top_k: int,
        alpha: float,
) -> list:
    """
    传统rag检索文本块
    content_augmentation: 文档增强类型， 1：无，2: 查找文本块的下一块
    """
    if retrieval_type != 1:
        # 获取关键词
        keywords = await get_keywords(query, llm)

    # 获取向量
    search_vector = await get_embedding(embed, query)

    # 检索
    if retrieval_type == 1:
        chunks = await Vectors.search_chunks_by_vector(collection_name=collection_name, query_vector=search_vector[0])
    elif retrieval_type == 2:
        chunks = await Vectors.search_chunks_by_bm25(collection_name=collection_name, query_keywords=keywords)
    else:
        chunks = await Vectors.search_chunks_by_hybrid(
            collection_name=collection_name,
            query_vector=search_vector[0],
            query_keywords=keywords,
            alpha=alpha,
        )

    # 文档增强
    if content_augmentation == 2:
        searched_chunks = chunks.copy()
        for chunk in chunks:
            response = Vectors.get_chunk_by_num(
                collection_name=collection_name,
                source=chunk.properties["source"],
                num=chunk.properties["num"]+1
            )
            if len(response) > 0 and response[0] in searched_chunks:
                continue
            searched_chunks.extend(response)
    else:
        searched_chunks = chunks



    # 重排
    contexts = [f"{x.properties['description']}\n{x.properties['text']}" for x in searched_chunks]
    scores = await get_rerank_scores(query=query, contexts=contexts, reranker=reranker)
    sorted_data = sorted(scores, key=lambda x: x['score'], reverse=True)
    # 获取top_k检索结果
    top_results = [searched_chunks[item['index']] for item in sorted_data[:top_k]]

    return top_results

